import torch
import numpy as np
from torch.utils.data import Dataset
from torch.utils.data import DataLoader
import matplotlib.pyplot as plt


# prepare dataset
class GetDataset(Dataset):
    def __init__(self, filepath):
        xy = np.loadtxt(filepath, delimiter=',', dtype=np.float32)
        self.len = xy.shape[0]  # shape(多少行，多少列)
        self.x_data = torch.from_numpy(xy[:, 1:])
        self.y_data = torch.from_numpy(xy[:, [0]])

    def __getitem__(self, index):
        return self.x_data[index], self.y_data[index]

    def __len__(self):
        return self.len
    
    def Get_Classes(self):
        classes=torch.unique(self.y_data)
        return classes.size(0)
    
    def Get_labels(self):
        return self.y_data

# data_set = GetDataset('.\data\ECG5000_TRAIN')
# train_loader = DataLoader(dataset=data_set, batch_size=10, shuffle=True)
# # for train_data in train_loader:
# train_data=list(train_loader)
# x = train_data[0]
# print(x[0].size(1))

